package com.training.utils;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.SecureRandom;
import java.util.Base64;

/**
 * AES加密工具类
 */
public class AESUtil {
    
    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/ECB/PKCS5Padding";
    /**
     * 从字符串加载密钥
     * @param key Base64编码的密钥字符串
     * @return SecretKey对象
     */
    public static SecretKey loadKey(String key) {
        // 如果输入是普通字符串，需要先转换为16字节的密钥
        if (key.length() < 16) {
            // 如果长度不足16位，用'0'填充
            key = String.format("%-16s", key).replace(' ', '0');
        } else if (key.length() > 16) {
            // 如果长度超过16位，截取前16位
            key = key.substring(0, 16);
        }
        
        return new SecretKeySpec(key.getBytes(), ALGORITHM);
    }
    
    /**
     * 加密字节数组
     * @param data 要加密的数据
     * @param key 密钥
     * @return 加密后的数据
     * @throws Exception 加密异常
     */
    public static byte[] encrypt(byte[] data, SecretKey key) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, key);
        return cipher.doFinal(data);
    }
    
    /**
     * 解密字节数组
     * @param encryptedData 要解密的数据
     * @param key 密钥
     * @return 解密后的数据
     * @throws Exception 解密异常
     */
    public static byte[] decrypt(byte[] encryptedData, SecretKey key) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, key);
        return cipher.doFinal(encryptedData);
    }
    
    /**
     * 加密文件
     * @param inputFile 输入文件
     * @param outputFile 输出文件
     * @param key 密钥
     * @throws Exception 加密异常
     */
    public static void encryptFile(File inputFile, File outputFile, SecretKey key) throws Exception {
        try (FileInputStream inputStream = new FileInputStream(inputFile);
             FileOutputStream outputStream = new FileOutputStream(outputFile)) {
            
            byte[] inputBytes = new byte[(int) inputFile.length()];
            inputStream.read(inputBytes);
            
            byte[] outputBytes = encrypt(inputBytes, key);
            
            outputStream.write(outputBytes);
        }
    }
    
    /**
     * 解密文件
     * @param inputFile 输入文件
     * @param outputFile 输出文件
     * @param key 密钥
     * @throws Exception 解密异常
     */
    public static void decryptFile(File inputFile, File outputFile, SecretKey key) throws Exception {
        try (FileInputStream inputStream = new FileInputStream(inputFile);
             FileOutputStream outputStream = new FileOutputStream(outputFile)) {
            
            byte[] inputBytes = new byte[(int) inputFile.length()];
            inputStream.read(inputBytes);
            
            byte[] outputBytes = decrypt(inputBytes, key);
            
            outputStream.write(outputBytes);
        }
    }
    
    /**
     * 检查密钥是否有效
     * @param key Base64编码的密钥字符串
     * @return 是否有效
     */
    public static boolean isValidKey(String key) {
        try {
            if (key == null || key.isEmpty()) {
                return false;
            }
            byte[] decodedKey = Base64.getDecoder().decode(key);
            // AES密钥长度必须是16、24或32字节
            return decodedKey.length == 16 || decodedKey.length == 24 || decodedKey.length == 32;
        } catch (Exception e) {
            return false;
        }
    }
}